import time
import torch
import argparse

from fvcore.nn import FlopCountAnalysis, parameter_count

from models.builder import MultiEncoderDecoder as rgbtcc


def parse_args():
    parser = argparse.ArgumentParser(description='Test')
    parser.add_argument('--data-dir', default='./data/processedDroneRGBT/',
                            help='training data directory')
    parser.add_argument('--save-dir', default='./pth_DroneRGBT/best_model.pth',
                            help='model directory')
    parser.add_argument('--dataset', default='DroneRGBT', choices=['RGBTCC', 'DroneRGBT'],
                        help='dataset name')
    parser.add_argument('--backbone', default='swin_b', choices=['swin_b'],
                        help='backbone name')
    parser.add_argument('--pretrained-model', default='',
                        help='the path of pretrained model')
    parser.add_argument('--bn-eps', type=float, default=1e-3,
                        help='batch normalization epsilon')
    parser.add_argument('--bn-momentum', type=float, default=0.1,
                        help='batch normalization momentum')
    parser.add_argument('--device', default='0', help='gpu device')
    args = parser.parse_args()
    return args


@torch.no_grad()
def get_flops_params(model, inputs):
    flops = FlopCountAnalysis(model, inputs)
    flops = flops.total() / 1000000000.

    params = parameter_count(model)
    params = params[""] / 1000000.

    return flops, params
    
    
if __name__ == '__main__':
    x = torch.randn(1, 3, 512, 512).cuda()
    inputs = [x, x]
    
    args = parse_args()
    model = rgbtcc(args)
    model = model.cuda()
    model.eval()
    
    flops, params = get_flops_params(model, inputs)
    
    print(f"[INFO] FVcore FLOPs(G)={round(flops, 2)}, FVcore Params(M)={round(params, 2)}")
    
    for _ in range(5):
        y = model(inputs)

    n = 100
    start_time = time.time()
    for _ in range(n):
        y = model(inputs)
    end_time = time.time()
    
    fps = n / (end_time - start_time)
    print(f"[INFO] FPS={round(fps, 2)}")
